import tenseal as ts
import numpy as np
import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
import os
import random
from faker import Faker
import matplotlib.pyplot as plt

# ==========================
# 动态安全参数配置（BFV专用）
# ==========================
def configure_security(security_level="balanced"):
    if security_level == "high":
        return 16384, 1032193  # 高安全级参数
    elif security_level == "balanced":
        return 8192, 1032193  # 平衡模式
    else:  # 'mobile'
        return 4096, 1032193  # 移动设备优化


# ==========================
# 计算可靠性检查
# ==========================
def check_reliability(enc_data, step_name):
    try:
        decrypted = enc_data.decrypt()
        print(f"{step_name} - 计算可靠，解密值: {decrypted[:2]}")
        return True
    except:
        print(f"{step_name} - 计算失败")
        return False


# ==========================
# 安全比较协议（使用BFV合法运算实现）
# ==========================
def secure_compare_bfv(a, b, k=10000):
    """
    BFV安全比较：a > b 则返回1，否则0
    使用BFV支持的运算模拟比较操作
    """
    diff = a - b  # 计算差值

    # 使用线性函数逼近符号函数
    sign_approx = diff * k

    return sign_approx


# ==========================
# 隐私保护聚合
# ==========================
def privacy_aggregate(enc_vectors):
    aggregated = enc_vectors[0].copy()
    for vec in enc_vectors[1:]:
        aggregated += vec
    return aggregated


# ==========================
# 生成随机医疗数据文件
# ==========================
def generate_medical_data(filename, num_records=100):
    """生成随机医疗数据并保存为CSV文件"""
    fake = Faker()
    data = {
        'name': [],
        'age': [],
        'glucose': [],
        'insulin': [],
        'label': []
    }

    for _ in range(num_records):
        age = random.randint(20, 80)
        glucose = round(random.uniform(70.0, 200.0), 1)
        insulin = round(random.uniform(4.0, 30.0), 1)

        # 基于年龄、血糖和胰岛素水平生成标签（糖尿病风险）
        risk_factor = (age - 50) * 0.02 + (glucose - 120) * 0.01 + (insulin - 15) * 0.02 + random.uniform(-0.5, 0.5)
        label = 1 if risk_factor > 0 else 0

        data['name'].append(fake.first_name())
        data['age'].append(age)
        data['glucose'].append(glucose)
        data['insulin'].append(insulin)
        data['label'].append(label)

    df = pd.DataFrame(data)
    df.to_csv(filename, index=False)
    print(f"已生成 {num_records} 条医疗数据并保存到 {filename}")
    return df


# ==========================
# 主程序
# ==========================
if __name__ == "__main__":
    # 设置数据文件路径
    data_file = "medical_data.csv"

    # 如果数据文件不存在，则生成新的随机数据
    if not os.path.exists(data_file):
        print(f"数据文件 {data_file} 不存在，正在生成随机医疗数据...")
        df = generate_medical_data(data_file, num_records=100)
    else:
        # 从CSV文件加载医疗数据
        print(f"\n从文件 {data_file} 加载医疗数据...")
        df = pd.read_csv(data_file)

    patients = df.to_dict('records')

    # 显示前5条记录和标签分布
    print("\n样本数据预览:")
    print(df.head())
    print(f"\n标签分布: 健康={sum(df['label'] == 0)}, 糖尿病={sum(df['label'] == 1)}")

    # ==========================
    # 初始化BFV环境
    # ==========================
    print("\n1. 初始化BFV环境...")
    security_level = "balanced"
    poly_deg, plain_mod = configure_security(security_level)

    context = ts.context(
        ts.SCHEME_TYPE.BFV,
        poly_modulus_degree=poly_deg,
        plain_modulus=plain_mod
    )
    context.generate_galois_keys()
    context.generate_relin_keys()

    print(f"安全配置: 多项式阶数={poly_deg}, 明文模数={plain_mod}")

    # ==========================
    # 医疗数据准备
    # ==========================
    # 提取特征
    ages = np.array([p["age"] for p in patients], dtype=np.float64)
    glucoses = np.array([p["glucose"] for p in patients], dtype=np.float64)
    insulins = np.array([p["insulin"] for p in patients], dtype=np.float64)
    labels = np.array([p["label"] for p in patients], dtype=np.float64)

    # 数据标准化并保存参数
    scaler = StandardScaler()
    X = np.c_[ages, glucoses, insulins]
    X_scaled = scaler.fit_transform(X)
    mean = scaler.mean_  # 保存特征均值
    scale = scaler.scale_  # 保存特征标准差

    # 使用逻辑回归模型
    print("\n训练逻辑回归模型...")
    model = LogisticRegression()
    model.fit(X_scaled, labels)

    # 增大量化因子提高精度
    scale_factor = 10000
    coef_int = (model.coef_[0] * scale_factor).astype(np.int64)
    intercept_int = (model.intercept_[0] * scale_factor).astype(np.int64)

    print(f"模型系数: 年龄={model.coef_[0][0]:.4f}, 血糖={model.coef_[0][1]:.4f}, 胰岛素={model.coef_[0][2]:.4f}")
    print(f"模型截距: {model.intercept_[0]:.4f}")
    print(f"数据集大小: {len(patients)} 条记录")

    # ==========================
    # 数据加密
    # ==========================
    print("\n2. 加密数据...")

    # 标准化特征后加密
    ages_scaled = (ages - mean[0]) / scale[0]
    glucoses_scaled = (glucoses - mean[1]) / scale[1]
    insulins_scaled = (insulins - mean[2]) / scale[2]

    enc_ages = ts.bfv_vector(context, ages_scaled)
    enc_glucoses = ts.bfv_vector(context, glucoses_scaled)
    enc_insulins = ts.bfv_vector(context, insulins_scaled)
    enc_labels = ts.bfv_vector(context, labels)

    check_reliability(enc_ages, "初始加密")

    # ==========================
    # 隐私计算
    # ==========================
    print("\n3. 执行隐私计算...")

    # 1. 平均年龄
    enc_age_sum = ts.bfv_vector(context, ages).sum()
    check_reliability(enc_age_sum, "年龄总和计算")

    # 2. 糖尿病患者比例
    enc_diabetic_count = enc_labels.sum()
    check_reliability(enc_diabetic_count, "患者计数")

    # 3. 安全比较（年龄>50，血糖>120）
    age_threshold = np.full(len(ages), 50, dtype=np.int64)
    glucose_threshold = np.full(len(glucoses), 120, dtype=np.int64)

    # 创建BFV格式的阈值
    age_threshold_bfv = ts.bfv_vector(context, age_threshold)
    glucose_threshold_bfv = ts.bfv_vector(context, glucose_threshold)

    # 安全比较
    enc_age_risk = secure_compare_bfv(ts.bfv_vector(context, ages), age_threshold_bfv)
    enc_glucose_risk = secure_compare_bfv(ts.bfv_vector(context, glucoses), glucose_threshold_bfv)

    # 解密比较结果
    age_risk_result = enc_age_risk.decrypt()
    glucose_risk_result = enc_glucose_risk.decrypt()

    # 计算高风险人数（年龄>50且血糖>120）
    high_risk_count = sum(1 for i in range(len(ages))
                          if age_risk_result[i] > 0 and glucose_risk_result[i] > 0)

    # 4. 健康评分
    enc_ages_orig = ts.bfv_vector(context, ages)
    enc_glucoses_orig = ts.bfv_vector(context, glucoses)
    enc_insulins_orig = ts.bfv_vector(context, insulins)
    enc_score = enc_ages_orig * 30 + enc_glucoses_orig * 50 + enc_insulins_orig * 20
    check_reliability(enc_score, "健康评分计算")

    # 5. 加密预测
    # 创建与加密向量维度相同的系数向量
    coef0_vector = np.full(len(ages), coef_int[0], dtype=np.int64)
    coef1_vector = np.full(len(glucoses), coef_int[1], dtype=np.int64)
    coef2_vector = np.full(len(insulins), coef_int[2], dtype=np.int64)
    intercept_vector = np.full(len(ages), intercept_int, dtype=np.int64)

    # 加密预测计算
    enc_pred = (enc_ages * coef0_vector +
                enc_glucoses * coef1_vector +
                enc_insulins * coef2_vector +
                intercept_vector)
    check_reliability(enc_pred, "风险预测计算")

    # ==========================
    # 解密与结果处理
    # ==========================
    print("\n4. 解密结果...")
    age_sum = enc_age_sum.decrypt()[0]
    avg_age = age_sum / len(ages)  # 在明文端计算平均

    diabetic_count = enc_diabetic_count.decrypt()[0]
    diabetic_percent = diabetic_count / len(ages) * 100  # 在明文端计算百分比

    # 处理健康评分
    enc_ages_orig = ts.bfv_vector(context, ages)
    enc_glucoses_orig = ts.bfv_vector(context, glucoses)
    enc_insulins_orig = ts.bfv_vector(context, insulins)
    scores = np.array((enc_ages_orig * 30 + enc_glucoses_orig * 50 + enc_insulins_orig * 20).decrypt()) / 100.0

    # 处理预测结果（logits转概率）
    preds_logits = np.array(enc_pred.decrypt()) / scale_factor
    preds = 1 / (1 + np.exp(-preds_logits))  # 应用sigmoid函数转换为概率

    # ==========================
    # 结果展示
    # ==========================
    print("\n===== 分析报告 =====")
    print(f"患者总数: {len(patients)}人")
    print(f"平均年龄: {avg_age:.1f}岁")
    print(f"糖尿病比例: {diabetic_percent:.1f}%")
    print(f"高风险人数(年龄>50且血糖>120): {high_risk_count}人")

    # 随机选择5个患者展示个性化结果
    print("\n随机5位患者个性化结果:")
    sample_indices = random.sample(range(len(patients)), min(5, len(patients)))
    for i in sample_indices:
        p = patients[i]
        print(f"{p['name']} (年龄:{p['age']}, 血糖:{p['glucose']}, 胰岛素:{p['insulin']}):")
        print(f"  健康评分: {scores[i]:.1f}")
        print(f"  糖尿病风险预测值: {preds[i]:.4f} ({'高风险' if preds[i] > 0.5 else '低风险'})")

    # ==========================
    # 验证（与明文对比）
    # ==========================
    plain_avg = np.mean(ages)
    plain_diabetic = (labels.sum() / len(labels)) * 100
    plain_high_risk = np.sum((ages > 50) & (glucoses > 120))
    plain_scores = (ages * 30 + glucoses * 50 + insulins * 20) / 100.0

    # 明文预测（使用logits和概率）
    plain_logits = model.decision_function(X_scaled)
    plain_probs = 1 / (1 + np.exp(-plain_logits))  # 明文概率

    # 计算各项误差
    avg_age_error = abs(avg_age - plain_avg)
    diabetic_percent_error = abs(diabetic_percent - plain_diabetic)
    high_risk_consistency = (high_risk_count == plain_high_risk)

    # 健康评分误差
    score_errors = np.abs(scores - plain_scores)
    max_score_error = np.max(score_errors)
    avg_score_error = np.mean(score_errors)

    # 风险预测误差（logits和概率）
    risk_errors_logits = np.abs(preds_logits - plain_logits)
    risk_errors_probs = np.abs(preds - plain_probs)
    max_risk_logits_error = np.max(risk_errors_logits)
    avg_risk_logits_error = np.mean(risk_errors_logits)
    max_risk_probs_error = np.max(risk_errors_probs)
    avg_risk_probs_error = np.mean(risk_errors_probs)

    # 统一格式输出
    print("\n===== 验证结果 =====")
    print(f"平均年龄: 明文{plain_avg:.2f} vs 加密{avg_age:.2f} | 误差: {avg_age_error:.4f}")
    print(f"糖尿病比例: 明文{plain_diabetic:.2f}% vs 加密{diabetic_percent:.2f}% | 误差: {diabetic_percent_error:.4f}%")
    print(f"高风险人数: 明文:{plain_high_risk}人 vs 加密：{high_risk_count}人 | 一致性: {high_risk_consistency}")
    print(f"健康评分最大误差: {max_score_error:.4f} | 平均误差: {avg_score_error:.4f}")
    print(f"风险预测logits最大误差: {max_risk_logits_error:.4f} | 平均误差: {avg_risk_logits_error:.4f}")
    print(f"风险预测概率最大误差: {max_risk_probs_error:.4f} | 平均误差: {avg_risk_probs_error:.4f}")

    # 实验结果判定
    if avg_risk_probs_error < 0.1 and max_risk_probs_error < 0.2 and high_risk_consistency:
        print("\n实验成功: 全同态加密计算保持高精度!")
    else:
        print("\n注意: 部分结果存在显著误差")

        # 可视化预测结果
    plt.figure(figsize=(12, 8))

    # 绘制加密预测概率
    plt.subplot(2, 1, 1)
    plt.scatter(range(len(preds)), preds, c='blue', alpha=0.6, label='Encrypted Probability')
    plt.scatter(range(len(plain_probs)), plain_probs, c='red', alpha=0.6, label='Plaintext Probability')
    plt.axhline(y=0.5, color='g', linestyle='--', label='Risk Threshold')
    plt.title('Diabetes Risk Prediction Probability Comparison')
    plt.xlabel('Patient ID')
    plt.ylabel('Predicted Probability')
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)

    # 绘制预测logits对比
    plt.subplot(2, 1, 2)
    plt.scatter(range(len(preds_logits)), preds_logits, c='blue', alpha=0.6, label='Encrypted Logits')
    plt.scatter(range(len(plain_logits)), plain_logits, c='red', alpha=0.6, label='Plaintext Logits')
    plt.axhline(y=0, color='g', linestyle='--', label='Logits Threshold')
    plt.title('Diabetes Risk Prediction Logits Comparison')
    plt.xlabel('Patient ID')
    plt.ylabel('Predicted Logits Value')
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)

    plt.tight_layout()
    plt.savefig('prediction_comparison.png')
    print("已保存预测对比图: prediction_comparison.png")